import numpy as np
import pandas as pd
import Env.EnvMain_Route as ENVM    # 该环境下无人机不会阵亡
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
import multiprocessing as mp
import os
import Route_select.MySQLFunc_11 as MSF11

os.environ["OMP_NUM_THREADS"] = "1"
MAX_EP=1000000

class DataGenerate1v1(mp.Process):
    def __init__(self,n_friend,n_enemy,EP,name,sql_name):
        super(DataGenerate1v1,self).__init__()
        self.N_F=n_friend
        self.N_E=n_enemy
        self.env=ENVM.ACEnv(n_friend,n_enemy)
        self.G_EP=EP
        self.name='P'+str(name)
        self.sql_name=sql_name

    def run(self):
        sql = MSF11.Mysql(self.name)
        while self.G_EP.value<MAX_EP:
            E,TE,index = 0.,9,0
            decisions = {'1': self.decision_1, '2': self.decision_2, '3': self.decision_3}
            strategy=self.First_step_strategy()
            for i in range(len(decisions)):
                self.env.reset()
                ob = self.env.FirstStep(strategy)  # 随机初始位置
                record=(ob[0][0][0],ob[0][0][1],ob[0][0][2],ob[1][0][0],ob[1][0][1],ob[1][0][2])
                if self.name=='P1':
                    self.env.move()
                    fig = plt.figure()
                    ax = fig.add_subplot(111, projection='3d')
                    self.env.render(fig, ax)
                for j in range(120):
                    ob = self.env.step(decisions[str(i+1)](ob))
                    if self.name=='P1':
                        self.env.render(fig, ax)
                    E += ob[0][2][0]    # 一回合的被摧毁期望
                if self.name=='P1':
                    plt.close()
                if E<TE:
                    TE=E
                    index=i
                E=0
                record+=(index,TE)
            sql.Insert_record('1v1_test_case',record)
            with mp.Lock():
                print('最合适的策略：',index,' E:',TE)
            with self.G_EP.get_lock():
                self.G_EP.value += 1

    def decision_1(self,ob):
        name='Go_stright'
        lable='route_1'
        angle=[[0,0]]*self.N_F
        attack=[[0]]*self.N_F

        return [angle,attack]

    def decision_2(self, ob):
        name = 'Go_left'
        lable = 'route_2'
        angle = [[0, 0.02]] * self.N_F
        attack = [[0]] * self.N_F

        return [angle, attack]

    def decision_3(self, ob):
        name = 'Go_right'
        lable = 'route_2'
        angle = [[0, -0.02]] * self.N_F
        attack = [[0]] * self.N_F

        return [angle, attack]

    def First_step_strategy(self):  # FirstStep的策略
        strategy = []
        select = np.random.randint(0, 120)
        strategy.append([0, select])
        return strategy

class DataGenerate2v2(mp.Process):
    def __init__(self,n_friend,n_enemy,EP,name,sql_name):
        super(DataGenerate2v2,self).__init__()
        self.N_F=n_friend
        self.N_E=n_enemy
        self.env=ENVM.ACEnv(n_friend,n_enemy)
        self.G_EP=EP
        self.name='P'+str(name)
        self.sql_name=sql_name

    def run(self):
        sql = MSF11.Mysql(self.name)
        record_buffer=[]    # 记录缓存，减少数据库交互堵塞
        count=1     # 回合计数
        while self.G_EP.value<MAX_EP:
            E,TE,index = 0.,19,0
            decisions = {'1': self.decision_1, '2': self.decision_2, '3': self.decision_3,
                         '4': self.decision_4, '5': self.decision_5, '6': self.decision_6,
                         '7': self.decision_7, '8': self.decision_8, '9': self.decision_9}
            strategy=self.First_step_strategy2()
            for i in range(len(decisions)):
                self.env.reset()
                ob = self.env.FirstStep(strategy)  # 随机初始位置
                record = (ob[0][0][0], ob[0][0][1], ob[0][0][2],ob[0][1][0], ob[0][1][1], ob[0][1][2],
                          ob[1][0][0], ob[1][0][1], ob[1][0][2],ob[1][1][0], ob[1][1][1], ob[1][1][2])
                if self.name=='P1':
                    self.env.move()
                    fig = plt.figure()
                    ax = fig.add_subplot(111, projection='3d')
                    self.env.render(fig, ax)
                for j in range(120):
                    D=decisions[str(i+1)](ob)
                    ob = self.env.step(D)
                    if self.name=='P1':
                        self.env.render(fig, ax)
                    E += ob[0][2][0]    # 一回合的被摧毁期望
                if self.name=='P1':
                    plt.close()
                if E<TE:
                    TE=E
                    index=i
                E=0
                record+=(index,TE)
            record_buffer.append(record)
            if count%10==0:
                sql.Insert_record2('2v2_test_case',record)
                record_buffer.clear()
            with mp.Lock():
                print('最合适的策略：',index,' E:',TE)
            with self.G_EP.get_lock():
                self.G_EP.value += 1
            count+=1

    def decision_1(self,ob):
        name='Go_stright'
        lable='route_1'
        angle=[[0,0]]*self.N_F
        attack=[[0]]*self.N_F
        return [angle,attack]

    def decision_2(self, ob):
        name = 'Go_left'
        lable = 'route_2'
        angle = [[0, 0.02]] * self.N_F
        attack = [[0]] * self.N_F
        return [angle, attack]

    def decision_3(self, ob):
        name = 'Go_right'
        lable = 'route_3'
        angle = [[0, -0.02]] * self.N_F
        attack = [[0]] * self.N_F
        return [angle, attack]

    def decision_4(self, ob):
        name = 'Go_right'
        lable = 'route_12'
        angle = [[0, 0],[0,0.02]]
        attack = [[0]] * self.N_F
        return [angle, attack]

    def decision_5(self, ob):
        name = 'Go_right'
        lable = 'route_13'
        angle = [[0, 0], [0,-0.02]]
        attack = [[0]] * self.N_F
        return [angle, attack]

    def decision_6(self, ob):
        name = 'Go_right'
        lable = 'route_21'
        angle = [[0, 0.02], [0,0]]
        attack = [[0]] * self.N_F
        return [angle, attack]

    def decision_7(self, ob):
        name = 'Go_right'
        lable = 'route_23'
        angle = [[0, 0.02], [0,-0.02]]
        attack = [[0]] * self.N_F
        return [angle, attack]

    def decision_8(self, ob):
        name = 'Go_right'
        lable = 'route_31'
        angle = [[0, -0.02], [0,0]]
        attack = [[0]] * self.N_F
        return [angle, attack]

    def decision_9(self, ob):
        name = 'Go_right'
        lable = 'route_32'
        angle = [[0, -0.02], [0,0.02]]
        attack = [[0]] * self.N_F
        return [angle, attack]

    def First_step_strategy2(self): # FirstStep的策略
        strategy=[]
        select1 = np.random.randint(0, 120)
        select2=np.random.randint(0,120)
        select_target=np.random.randint(0,2,2)
        strategy.append([select_target[0],select1])
        strategy.append([select_target[1],select2])
        return strategy


if __name__=='__main__':
    EP = mp.Value('i', 0)
    workers = [DataGenerate2v2(2,2,EP,i,'1V1_test_case') for i in range(mp.cpu_count())]  #mp.cpu_count()
    [w.start() for w in workers]
    [w.join() for w in workers]

